package gossip.v2.message.serializer;

import com.baidu.bjf.remoting.protobuf.Codec;
import com.baidu.bjf.remoting.protobuf.ProtobufProxy;
import gossip.v2.message.*;
import gossip.v2.util.SerializationUtils;
import java.io.*;
import java.util.EnumMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class GossipMessageSerializer {
    private static final Map<Class<? extends GossipMessage>, Codec<? extends GossipMessage>> CODEC_CACHE = new ConcurrentHashMap<>();
    private static final Map<GossipMessage.Type, Class<? extends GossipMessage>> TYPE_TO_CLASS = new EnumMap<>(GossipMessage.Type.class);

    static {
        registerTypeMapping(HeartbeatMessage.class, GossipMessage.Type.HEARTBEAT);
        registerTypeMapping(SyncMessage.class, GossipMessage.Type.SYNC);
        registerTypeMapping(AckMessage.class, GossipMessage.Type.ACK);
        registerTypeMapping(UserMessage.class, GossipMessage.Type.USER);
        registerTypeMapping(HandshakeMessage.class, GossipMessage.Type.HANDSHAKE);
        registerTypeMapping(IndirectCheckRequest.class, GossipMessage.Type.INDIRECT_CHECK);
    }

    private static <T extends GossipMessage> void registerTypeMapping(Class<T> clazz, GossipMessage.Type type) {
        try {
            CODEC_CACHE.put(clazz, ProtobufProxy.create(clazz));
            TYPE_TO_CLASS.put(type, clazz);
        } catch (Exception e) {
            throw new RuntimeException("Failed to create codec for " + clazz.getName(), e);
        }
    }

    public byte[] serialize(GossipMessage message) {
        try {
            GossipMessage.Type type = message.getType();
            if (type == null) {
                throw new IOException("Message type is null: " + message.getClass());
            }

            @SuppressWarnings("unchecked")
            Codec<GossipMessage> codec = (Codec<GossipMessage>) CODEC_CACHE.get(message.getClass());
            if (codec == null) {
                throw new IOException("No codec registered for class: " + message.getClass());
            }

            byte[] bodyBytes = codec.encode(message);

            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            DataOutputStream dos = new DataOutputStream(bos);
            dos.writeByte(type.ordinal());
            dos.write(bodyBytes);
            return bos.toByteArray();
        } catch (IOException e) {
            throw new SerializationUtils.SerializationException("Serialization failed", e);
        }
    }

    public GossipMessage deserialize(byte[] data) {
        try {
            ByteArrayInputStream bis = new ByteArrayInputStream(data);
            DataInputStream dis = new DataInputStream(bis);

            // 1. 读取类型标识符
            int typeOrdinal = dis.readByte() & 0xFF;
            if (typeOrdinal >= GossipMessage.Type.values().length) {
                throw new IOException("Invalid message type: " + typeOrdinal);
            }

            // 2. 获取消息类型
            GossipMessage.Type type = GossipMessage.Type.values()[typeOrdinal];

            // 3. 获取对应消息类
            Class<? extends GossipMessage> messageClass = TYPE_TO_CLASS.get(type);
            if (messageClass == null) {
                throw new IOException("Unsupported message type: " + type);
            }

            // 4. 获取编解码器
            Codec<? extends GossipMessage> codec = CODEC_CACHE.get(messageClass);
            if (codec == null) {
                throw new IOException("No codec for type: " + type);
            }

            // 5. 读取消息体并反序列化
            byte[] body = new byte[data.length - 1];
            dis.readFully(body);
            return codec.decode(body);
        } catch (IOException e) {
            throw new SerializationUtils.SerializationException("Deserialization failed", e);
        }
    }
}
